package com.gper.edu.flink.stream.process;

import com.gper.edu.flink.stream.StreamEnvironment;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.MapState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.streaming.api.functions.ProcessFunction;
import org.apache.flink.streaming.api.functions.timestamps.BoundedOutOfOrdernessTimestampExtractor;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.evictors.TimeEvictor;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.triggers.ContinuousEventTimeTrigger;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.time.Duration;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

/**
 * @Author: guan.xianchun
 * @Description:
 * @Date: 2020/11/26 14:31
 * @Version: 1.0
 */
public class ProcessWindowFunctionDemo extends StreamEnvironment {

    public static void main(String[] args) {
        new ProcessWindowFunctionDemo().runJob();
    }

    @Override
    public void execute() throws Exception {
        setParallelism(1);
        DataStreamSource<String> source = getDataStreamSource();
        DataStream<Tuple2<String, Long>> watermarkStream = source.flatMap(new FlatMapFunction<String, Tuple2<String, Long>>() {
            @Override
            public void flatMap(String value, Collector<Tuple2<String, Long>> out) throws Exception {
                String[] split = value.split(",");
                out.collect(Tuple2.of(split[0], Long.valueOf(split[1])));
            }
        }).assignTimestampsAndWatermarks(WatermarkStrategy.<Tuple2<String, Long>>forBoundedOutOfOrderness(Duration.ofSeconds(3))
                .withTimestampAssigner(new SerializableTimestampAssigner<Tuple2<String, Long>>() {
                    @Override
                    public long extractTimestamp(Tuple2<String, Long> element, long recordTimestamp) {
                        return element.f1;
                    }
                }));
        watermarkStream
                .keyBy(new KeySelector<Tuple2<String, Long>, String>() {
                    @Override
                    public String getKey(Tuple2<String, Long> value) throws Exception {
                        return value.f0;
                    }
                })
//        .window(TumblingEventTimeWindows.of(Time.minutes(1)))
//        .trigger(ContinuousEventTimeTrigger.of(Time.seconds(10)))
//        .evictor(TimeEvictor.of(Time.seconds(0),true))
                .process(new KeyedProcessFunction<String, Tuple2<String, Long>, Tuple2<String, Long>>() {
                    private ListState<Tuple2<String, Long>> dataListState;
                    private MapState<Long, ArrayList<Tuple2<String, Long>>> mapState;
                    private Boolean userMapState = false;

                    @Override
                    public void open(Configuration parameters) throws Exception {
                        dataListState = getRuntimeContext().getListState(new ListStateDescriptor<Tuple2<String, Long>>("_list_state",
                                TypeInformation.of(new TypeHint<Tuple2<String, Long>>() {
                                })));
                        mapState = getRuntimeContext().getMapState(new MapStateDescriptor<Long, ArrayList<Tuple2<String, Long>>>("_map_state",
                                TypeInformation.of(Long.class), TypeInformation.of(new TypeHint<ArrayList<Tuple2<String, Long>>>() {
                        })));
                    }

                    private void addElement(Tuple2<String, Long> value) throws Exception {
                        if (userMapState) {
                            if (!mapState.contains(value.f1)) {
                                mapState.put(value.f1, new ArrayList<>());
                            }
                            ArrayList<Tuple2<String, Long>> list = mapState.get(value.f1);
                            list.add(value);
                            mapState.put(value.f1, list);
                        } else {
                            dataListState.add(value);
                        }
                    }

                    @Override
                    public void processElement(Tuple2<String, Long> value, Context ctx, Collector<Tuple2<String, Long>> out) throws Exception {
                        System.out.println(Thread.currentThread().getId() + " add  " + value.toString());
                        addElement(value);
                        ctx.timerService().registerEventTimeTimer(value.f1);
                    }

                    protected void listStateOutput(OnTimerContext ctx, Collector<Tuple2<String, Long>> out) throws Exception {
                        Iterator<Tuple2<String, Long>> iterator = dataListState.get().iterator();
                        long cnt = 0;
                        while (iterator.hasNext()) {
                            cnt++;
                            System.out.println("data in listState :" + iterator.next());
//                            iterator.remove();  //ListSate无法通过iterator.remove来删除状态中的数据，必须通过.update方法
                        }
                        out.collect(Tuple2.of(ctx.getCurrentKey(), cnt));
                        dataListState.update(new ArrayList<>());
                    }

                    protected void mapStateOutput(OnTimerContext ctx, Collector<Tuple2<String, Long>> out) throws Exception {
                        Iterator<Long> iterator = mapState.keys().iterator();
                        long cnt = 0;
                        while (iterator.hasNext()) {
                            Long key = iterator.next();
                            for (Tuple2 data : mapState.get(key)) {
                                cnt++;
                                System.out.println("data in mapState : " + data.toString());
                            }
                            iterator.remove();  //MapSate可以通过iterator.remove来删除状态中的数据，或通过.remove方法删除key对应的数据
                        }
                        out.collect(Tuple2.of(ctx.getCurrentKey(), cnt));
                    }

                    @Override
                    public void onTimer(long timestamp, OnTimerContext ctx, Collector<Tuple2<String, Long>> out) throws Exception {
                        System.out.println(Thread.currentThread().getId() + " " + timestamp);
                        if (userMapState) {
                            mapStateOutput(ctx, out);
                        } else {
                            listStateOutput(ctx, out);
                        }
                    }
                }).print("======");
    }
}
